import numpy as np  # 导入numpy用于数学运算
import matplotlib.pyplot as plt  # 导入matplotlib用于绘图
from random import sample, random  # 导入随机模块用于随机选择
from operator import itemgetter  # 导入itemgetter用于排序时快速获取元素


class TSPSolver:  # 定义解决旅行商问题(TSP)的类
    def __init__(self, filename):  # 构造函数
        self.cities = self.load_cities(filename)  # 加载城市坐标
        self.distance_cache = {}  # 缓存距离计算结果

    # 读取文件中的城市坐标数据
    def load_cities(self, filename):
        cities = {}
        with open(filename, 'r') as file:  # 打开文件并逐行读取
            for line in file:
                city_id, x, y = map(int, line.split())  # 解析每一行的数据
                cities[city_id] = (x, y)  # 将城市坐标存储到字典中
        return cities

    # 计算两点之间的欧氏距离，并缓存结果
    def distance(self, city1, city2):
        if (city1, city2) in self.distance_cache:
            return self.distance_cache[(city1, city2)]
        dist = ((city1[0] - city2[0]) ** 2 + (city1[1] - city2[1]) ** 2) ** 0.5  # 计算距离
        self.distance_cache[(city1, city2)] = dist  # 存储距离到缓存
        return dist

    # 计算给定路径的总长度
    def path_length(self, path):
        length = 0
        n = len(path)
        for i in range(n):
            j = (i + 1) % n  # 下一个城市的索引，如果到达最后一个则返回第一个
            length += self.distance(self.cities[path[i]], self.cities[path[j]])  # 累加距离
        return length

    # 生成初始种群，每个个体都是城市的一个排列
    def generate_population(self, size):
        city_ids = list(self.cities.keys())  # 获取城市ID列表
        return [sample(city_ids, len(city_ids)) for _ in range(size)]  # 创建指定数量的个体

    # 使用选择机制选出优秀个体进入下一代
    def selection(self, population, fitnesses, elite_size):
        sorted_pop_fit = sorted(zip(population, fitnesses), key=itemgetter(1))  # 按适应度排序
        selected = [ind for ind, fit in sorted_pop_fit[:elite_size]]  # 选取前elite_size个个体
        return selected

    # 进行有序交叉操作生成新个体
    def ordered_crossover(self, parent1, parent2):
        size = len(parent1)
        pos1, pos2 = sorted(sample(range(size), 2))  # 随机选取交叉点

        child = [-1] * size  # 初始化子代
        child[pos1:pos2] = parent1[pos1:pos2]  # 复制父代片段到子代

        index = pos2
        for element in parent2:
            if element not in child:
                while child[index % size] != -1:
                    index = (index + 1) % size
                child[index % size] = element
                index = (index + 1) % size
        return child

    # 对个体进行变异操作
    def mutate(self, individual):
        pos1, pos2 = sorted(sample(range(len(individual)), 2))  # 随机选择两个位置交换
        individual[pos1], individual[pos2] = individual[pos2], individual[pos1]

    # 主函数，使用遗传算法求解TSP问题
    def solve_tsp(self, population_size, generations, mutation_rate, elite_size):
        population = self.generate_population(population_size)  # 生成初始种群

        best_fitness_history = []  # 保存每一代的最佳路径长度

        for gen in range(generations):
            fitnesses = [self.path_length(ind) for ind in population]  # 计算种群适应度
            best_fitness_history.append(min(fitnesses))  # 记录当前代的最佳适应度
            elite_population = self.selection(population, fitnesses, elite_size)  # 选择精英个体

            children = []
            while len(children) < population_size - elite_size:
                parent1, parent2 = sample(elite_population, 2)  # 从精英个体中随机选择父母
                child = self.ordered_crossover(parent1, parent2)  # 交叉产生子代
                if random() < mutation_rate:
                    self.mutate(child)  # 变异
                children.append(child)  # 添加到后代中

            population = elite_population + children  # 更新种群

        best_path = min(population, key=lambda path: self.path_length(path))  # 寻找最佳路径
        return best_path, best_fitness_history


# 主函数
def main():
    filename = r'C:\Users\86191\Desktop\资料PPT\人工智能\city_data.txt'  # 文件路径
    tsp_solver = TSPSolver(filename)  # 创建TSP求解器实例
    population_size = 140
    generations = 1800
    mutation_rate = 0.4198
    elite_size = 60

    best_path, best_fitness_history = tsp_solver.solve_tsp(population_size, generations, mutation_rate, elite_size)
    print(f"Best path found: {best_path}")
    print(f"Path length: {tsp_solver.path_length(best_path)}")

    # 可视化最佳路径
    x, y = zip(*[tsp_solver.cities[city] for city in best_path + [best_path[0]]])
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 2, 1)
    plt.plot(x, y, 'o-', markersize=5, label='Path')

    # 添加城市编号和箭头
    for idx, city in enumerate(best_path):
        plt.text(x[idx], y[idx], str(city), fontsize=9, ha='right')
        if idx == len(best_path) - 1:
            plt.annotate('', xy=(x[0], y[0]), xytext=(x[idx], y[idx]),
                         arrowprops=dict(arrowstyle="->", lw=0.5, color='red'))
        else:
            plt.annotate('', xy=(x[idx + 1], y[idx + 1]), xytext=(x[idx], y[idx]),
                         arrowprops=dict(arrowstyle="->", lw=0.5, color='red'))

    plt.title('Best TSP Path')
    plt.legend()

    # 绘制路径长度变化趋势
    plt.subplot(1, 2, 2)
    plt.plot(best_fitness_history, label='Best Fitness', color='blue')
    plt.xlabel('Generation')
    plt.ylabel('Best Path Length')
    plt.title('Fitness Trend')
    plt.legend()

    plt.tight_layout()
    plt.show()

if __name__ == "__main__":
    main()